#!/usr/local/bin/python
import resource,copy
from Bio.Blast import NCBIXML
from Bio import SeqIO
from ProjectBio import ParseDefline, nr_dict, ParseTC
from Bio.Blast.Applications import NcbiblastpCommandline
from urllib import quote as urlencode
from urllib2 import urlopen
import hmmtop,shutil
from math import ceil
import sys,os,pickle
import matplotlib
import hmmgap
import tcdb
import cdd
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pylab,re,tempfile
rsrc = resource.RLIMIT_DATA
soft, hard = resource.getrlimit(rsrc)
resource.setrlimit(rsrc, (1073741824, hard)) #limit to one gig, omg..

class Tools:
	
	def __init__(self):
		self.dbfile = '/db/tcdb'
		self.tcdb = os.environ['HOME']+self.dbfile
		self.subject = False
		self.indir = False
		self.goodresults = []
		self.bestresults = []
		self.notmsresults = []
		self.debug=True
		self.tms = {}
		self.expect = 0.001
		self.minlen = 50
		self.mysubjects = False
		self.mytcdb = False
		self.rows = []
		self.data = []
		self.globalcount = 0
		self.cdd_on = False
		self.abbreviations = {}
		self.ortho = False
		self.subject_gis = False
		self.target_gis = False
		self.alabel='Query'
		self.blabel='TCDB-Hit'
		self.names = tcdb.Names()
		self.substrates = tcdb.Substrates()
		tcdb.use_local()
		tcdb.use_local_betabarrel()
		
	def blast_all(self):
		try:
			os.makedirs(self.indir+"/xml")
		except:
			pass
		if self.ortho is not False:
			self.prep_orthologs()
		subjects = SeqIO.parse(open(self.subject),'fasta')
		for subject in subjects:
			subject_file = tempfile.NamedTemporaryFile()
			SeqIO.write(subject,subject_file,'fasta')
			subject_file.flush()
			blast_out = "'"+self.indir+'/xml/'+subject.id+".xml"+"'"
			if os.path.exists(blast_out[1:-1]) is True: # removes the quotes in a rly cool way :)
				continue # Blast xml already exists...
			blastp = NcbiblastpCommandline(query=subject_file.name, db=self.tcdb, evalue=self.expect, out=blast_out, outfmt=5)
			blastp()
			print "Blasted :: %s" %subject.id
			
	def prep_orthologs(self):
		subjects = SeqIO.parse(open(self.subject),'fasta')
		targets = SeqIO.parse(open(self.ortho),'fasta')
		# Load subject Gis
		subject_gis = open(self.subject_gis,'r')
		target_gis = open(self.target_gis,'r')
		subjects = SeqIO.to_dict(subjects)
		targets = SeqIO.to_dict(targets)
		# Make ortho directory
		try:
			os.makedirs(self.indir+"/orthologs")
		except:
			pass
		# Write subjects & targets
		mysubjects = open(self.indir+'/orthologs/mysubjects.faa','wb')
		mytargets = open(self.indir+'/orthologs/mytargets.faa','wb')
		for sgi in subject_gis:
			SeqIO.write(subjects[sgi.strip()],mysubjects,'fasta')
		for tgi in target_gis:
			SeqIO.write(targets[tgi.strip()],mytargets,'fasta')
		mysubjects.flush()
		mytargets.flush()
		os.system('makeblastdb -in '+mytargets.name)
		self.subject = mysubjects.name
		self.tcdb = mytargets.name
			
	
	def load_good(self):
		db = self.indir+"/goodresults.db"
		if os.path.exists(db) and self.debug:
			self.goodresults = pickle.load(open(db,'r'))
			return
		xml = os.listdir(self.indir+"/xml")
		rez = []
		for res in xml:
			if os.path.exists(self.indir+'/xml/'+res) is False:
				continue
			try:
				for blast in NCBIXML.parse(open(self.indir+'/xml/'+res)):
					subject =  blast.query
					try:
						descriptions = [i for i in blast.descriptions]
						alignments = [i for i in blast.alignments]
						results = zip(descriptions,alignments)
						results.sort(key=lambda x:x[0].e,reverse=False)
						record = results[0] # Contains top description and top alignment object in tuple.
						hsps = [i for i in record[1].hsps]
						hsps.sort(key=lambda x:x.expect)
						hsp = hsps[0]
						if record[0].e > self.expect or len(hsp.match) < self.minlen:
							continue
						rez.append((subject,record,hsp)) # (genome ID, hit record <e,title>, hit.hsp)
					except:
						pass
			except:
				continue
		self.goodresults = rez
		pickle.dump(self.goodresults,open(db,'wb'))
		return
		
	def write_fastas(self):
		mytcdb = []
		mysubject = []
		tcdb = SeqIO.parse(self.tcdb,'fasta')
		tcdb=nr_dict(tcdb)
		subjects= SeqIO.parse(self.subject,'fasta')
		subjects= SeqIO.to_dict(subjects)
		for subject,hit,hsp in self.goodresults:
			hit = hit[0]
			subject=ParseDefline(subject).id
			mysubject.append(subjects[str(subject)])
			try:
				mytcdb.append(tcdb[ParseDefline(hit.title,True).id])
			except:
				(family,tcid,acc) = ParseTC(hit.title)
				#print tcid,acc
				print ParseDefline(hit.title,True).id
				#print hit.title
				quit()
				
		subject_file=open(self.indir+"/mysubjects.faa",'wb')
		tcdb_file=open(self.indir+"/mytcdb.faa",'wb')
		SeqIO.write(list(set(mysubject)),subject_file,'fasta')
		SeqIO.write(list(set(mytcdb)),tcdb_file,'fasta')
	
	def hmmtop(self):
		db = self.indir+'/hmmtop.db'
		if os.path.exists(db) and self.debug:
			self.tms = pickle.load(open(db,'r'))
			return
		ht = hmmtop.tools()
		ht.add_library('subjects',self.indir+"/mysubjects.faa") # Genome
		ht.add_library('tcdb',self.indir+"/mytcdb.faa")
		ht.scan_libraries()
		pickle.dump(ht.results,open(db,'wb'))
		self.tms = ht.results
		return
		
	def calculate_tms_scores(self):
		for genome,tcdb,hsp in self.goodresults:
			genome=ParseDefline(genome).id
			tcdb=tcdb[0]
			delta  = hsp.sbjct_start-hsp.query_start
			tcdbid = ParseDefline(tcdb.title,True).id
			try:
				genome_tms = self.tms['subjects'][genome].values()
				tcdb_tms = self.tms['tcdb'][tcdbid].values()
				
			except KeyError:
				self.notmsresults.append((genome,tcdb,hsp,None)) #Genome or TCDB hit dont have a TMS. These must be manually revised later!
				continue
			g_tms = [[i[0]+delta,i[1]+delta] for i in genome_tms]
			overlap = self.find_overlap(g_tms,tcdb_tms)
			row= (genome,tcdb,hsp,overlap)
			self.bestresults.append(row)
		self.bestresults.sort(key=lambda x:(x[1].e,x[3],self.tcsort(x[1].title)))
		self.notmsresults.sort(key=lambda x:(x[1].e,self.tcsort(x[1].title)))

		
	
	def find_overlap(self,subject,target):
		subjects = [range(i[0],i[1]+1) for i in subject]
		targets = [range(i[0],i[1]+1) for i in target]
		overlap = []
		for sub in subjects:
			for tar in targets:
				overlap.extend(set(sub)&set(tar))
		return float(len(set(overlap)))/20
		
	def title_extract(self,string): #returns: (acc,tcid)
		string = string.split(" ")
		gi = string[0].split('|')[-1]
		tcid = string[1]
		return (gi,tcid)
		
	def tcsort(self,title):
		if self.ortho is not False:
			return 1
		title=ParseDefline(title,True).description
		(family,tc,acc) = ParseTC(title)
		tc = tc.split('.')
		return ( int(tc[0]), str(tc[1]), int(tc[2]), int(tc[3]), int(tc[4]) )
	
	def build_view(self,data):
		(genome,tcdb,hsp,overlap) = data
		try:
			os.mkdir(self.indir+"/img")
		except:
			pass
		genome=ParseDefline(genome).id
		tid = ParseDefline(tcdb.title,True).id
		if os.path.exists(self.indir+"/img/"+genome+".png") is False:
			try:
				san = self.subjecthmg(self.mysubjects[genome],hsp.query)
				tan = self.tcdbhmg(self.mytcdb[tid],hsp.sbjct)
				self.what(hsp.query,hsp.sbjct,self.indir+"/img/"+genome+".png",[san,tan])
			except:
				print hsp.query
				print hsp.sbjct
				print genome
				print 'error, quit'
				quit()
		
		(family,tcid,acc) = ParseTC(ParseDefline(tcdb.title,True).description)
		try:
			query_tms = len(self.tms['subjects'][genome])
		except:
			query_tms = 0
		try:
			hit_tms = len(self.tms['tcdb'][ParseDefline(tcdb.title,True).id])
		except:
			hit_tms = 0
		self.globalcount += 1
		if self.ortho is False:
			family = self.names.get_family_abr('.'.join(tcid.split('.')[0:3]))
		else:
			family = 'family_place_holder'
		glink = '<a href="content.html#%s">%s</a>'%(genome,genome)
		ident = round((float(hsp.identities)/len(hsp.match))*100)
		mysubstrate = self.substrates.get_tcid_substrates(tcid)
		mysubstrate = ", ".join(mysubstrate) if mysubstrate is not None else None
		row = (glink,acc,tcid,tcdb.title,len(hsp.match),tcdb.e,ident,query_tms,\
		hit_tms,overlap,family,mysubstrate,self.globalcount)
		row =['<td>'+str(i)+'</td>' for i in row]
		htmlrow = "<tr>\n%s\n</tr>"%("\n\t".join(row))
		#self.rows.append(htmlrow)
		if self.cdd_on is True:
			mycdd = 'Query & TC-Hit Conserved Domains:<br><img src=\'./cdd/%s.1.png\'><br><img src=\'./cdd/%s.2.png\'><br>'%(genome,genome)
		else:
			mycdd = ''
		ol = float(overlap) if overlap is not None else float(0)
		content = "<div class='result' id='%s'> <h3><a name='%s'>%s</a></h3>  <p>Hit Gi: %s<br>   Hit TCID: %s</p> <p>Hit Description: %s<br>   \
		<br>   Mach Len: %i<br>   e:%f</p> <p>Query TMS Count : %i<br>   Hit TMS Count: %i     <br>     TMS-Overlap Score: %f<br>     \
		Predicted Substrates:%s <br><br>     BLAST Alignment:<br> 	<pre> 	%s 	</pre> <br>  <br> Binary-Hydropathy Alignment View:<br> 	\
		<img src='./img/%s.png'><br>%s </p> </div>" %(genome,genome,genome,acc,tcid,tcdb.title,\
		len(hsp.match),tcdb.e,query_tms,hit_tms,ol,mysubstrate,str(hsp),urlencode(genome),mycdd)
		#self.data.append(content)
		return htmlrow,content
		
	def write_results(self):
		bestresults = copy.deepcopy(self.bestresults)
		bestresults.extend(self.notmsresults)
		results = open(self.indir+'/results.html','wb')
		content = open(self.indir+'/content.html','wb')
		html = '<html><table width="100%%" border="1"> <tr> <td>Query Gi</td> <td>Hit Gi</td> \
		<td>Hit TCID</td> <td>Hit Desc</td> <td>Match Len</td> <td>e-Val</td> <td>% Identity</td> <td>Query TMS#</td>\
		<td>Hit TMS#</td> <td>TM-Overlap Score</td> <td>Family Abrv.</td><td>Predicted Substrate</td> <td> #</td> </tr><br>\n\n'
		results.write(html)
		content.write("<html>")
		self.mysubjects = SeqIO.parse(self.indir+'/mysubjects.faa','fasta')
		self.mytcdb = SeqIO.parse(self.indir+'/mytcdb.faa','fasta')
		self.mysubjects = SeqIO.to_dict(self.mysubjects)
		self.mytcdb = SeqIO.to_dict(self.mytcdb)
		if self.cdd_on:
			self.cdd_extract()
		self.subjecthmg = hmmgap.annotate()
		self.tcdbhmg = hmmgap.annotate()
		self.subjecthmg.hmmtop = self.tms['subjects']
		self.tcdbhmg.hmmtop = self.tms['tcdb']
		for res in bestresults:
			(row,data) = self.build_view(res)
			results.write(row)
			content.write(data)
			print "Generated Results for :: %s" %ParseDefline(res[0]).id
		results.write("</table></html>")
		content.write("</html>")
		
	def cdd_extract(self):
		if os.path.exists(self.indir+'/cdd') is False:
			os.mkdir(self.indir+'/cdd')
		fastas = dict(self.mysubjects,**self.mytcdb)
		thisdir = self.indir+'/cdd/'
		for subject,hit,hsp in self.goodresults:
			if os.path.exists(thisdir+subject.title()+'.1.png') is False:
				cdd.fetch(str(hsp.query),thisdir+subject.title()+'.1.png')
			if os.path.exists(thisdir+subject.title()+'.2.png') is False:
				cdd.fetch(str(hsp.sbjct),thisdir+subject.title()+'.2.png')
	
	def hydro(self,gseq):
		seq=gseq.replace('-','')
		window = 19
		prev = 0
		index={'G':(-0.400,0.48),'I':(4.500,1.38),'S':(-0.800,-0.18),'Q':(-3.500,-0.85),'E':(-3.500,-0.74),\
		'A':(1.800,0.62),'M':(1.900,0.64),'T':(-0.700,-0.05),'Y':(-1.300,0.26),'H':(-3.200,-0.4),\
		'V':(4.200,1.08),'F':(2.800,1.19),'C':(2.500,0.29),'W':(-0.900,0.81),'K':(-3.900,-1.5),\
		'L':(3.800,1.06),'P':(-1.600,0.12),'N':(-3.500,-0.78),'D':(-3.500,-0.90),'R':(-4.500,-2.53),'X':(0,0),'U':(0,0)}
		midpt = (window+1)/2
		length = len(seq)
		hydro = []
		for i in range(length-window+1):
			total = 0
			for j in range(window):
				total +=index[seq[i+j]][0]
			total = total/window
			hydro.append(total)
		if len(seq) == len(gseq):
			return hydro
		replace = re.finditer('(-+)',gseq)
		inserts = {}
		for i in replace:
			inserts.setdefault(i.start(),i.end()-i.start())
		first = False
		newhydro = []
		for x, h in enumerate(hydro):
			if x in inserts.keys() and first is False:
				first = True
				for y in range(inserts[x]):
					newhydro.append(0)
				newcount = x + inserts[x]
				continue
			if first is False:
				newhydro.append(h)
				continue
			if first is True and newcount in inserts.keys():
				for y in range(inserts[newcount]):
					newhydro.append(0)
				newcount += inserts[newcount]
				continue
			else:
				newhydro.append(h)
				newcount +=1

		return newhydro
		
	def what(self,a,b,outfile,hmt=[]):
		#quit()
		ha = self.hydro(a)
		hb = self.hydro(b)
		omg = [len(ha),len(hb)]
		readbar = re.compile(r'\d+?[+_]+\+')
		omg.sort()
		ha =ha[0:omg[0]]
		hb =hb[0:omg[0]]
		x_data=range(0,len(ha))
		plt.figure()
		plt.axhline(y=0,color='black')
		plt.ylim(-3,3)
		plt.xlim(right=len(a))
		plt.plot(x_data,ha,linewidth=1,label=self.alabel,color='blue')
		plt.plot(x_data,hb,linewidth=1,label=self.blabel,color='red')
		plt.xlabel("Residue #")
		plt.ylabel("Hydro")
		plt.legend(loc='lower right')
		# Draw TMS bars
		if len(hmt) == 2:
			sub = readbar.finditer(str(hmt[0]))
			tar = readbar.finditer(str(hmt[1]))
			for tms in sub:
				subtract = (len(hmt[0])-len(ha))/2
				if tms.end()-subtract>len(ha):
					subtract = tms.end()-len(ha)
				plt.axvspan(tms.start()-subtract,tms.end()-subtract, facecolor="blue", alpha=0.3)
			for tms in tar:
				subtract = (len(hmt[1])-len(ha))/2
				if tms.end()-subtract>len(hb):
					subtract = tms.end()-len(hb)
				plt.axvspan(tms.start()-subtract,tms.end()-subtract, facecolor="red", alpha=0.3)
		fig = matplotlib.pyplot.gcf()
		fig.set_size_inches(15,3)
		plt.savefig(outfile, dpi=80, format="png",bbox_inches='tight', pad_inches=0.003)
		plt.clf()
		
if __name__=="__main__":
	from optparse import OptionParser,OptionGroup
	desc="Welcome to GBlast! Easily identify transporters in entire Genomes/Proteomes - By Vamsee Reddy"
	version = "GBlast V3"
	opts = OptionParser(description=desc,version=version)
	opts.add_option('-i',action='store',type='string',dest='input', help="Path to genome/proteome file")
	opts.add_option('-o',action='store',type='string',dest='output',default='genome.fsa', help="Results output name")
	opts.add_option('--evalue',action='store',type='float',dest='evalue',default=0.001, help="Minimum e-Value [0.001]")
	opts.add_option('--betabarrel',action='store_true', dest='bb',default=False, help="Find Beta Barrels instead of TMS")
	opts.add_option('--cdd',action='store_true',default=False, dest='cdd', help="Include CDD analysis (Takes a while)")
	opts.add_option('--orthologs',action='store', dest='ortho',default=False, help="Find orthologs with this fasta file (optional)")
	opts.add_option('--subject_gi',action='store', dest='subgi',default=False, help="Ortholog search using list of only these subjects")
	opts.add_option('--target_gi',action='store', dest='targi',default=False, help="Ortholog search using list of only these targets")
	(cli,args)=opts.parse_args()
	if cli.input is not None and cli.output is not None:
		GB = Tools()
		if(cli.bb):
			GB.dbfile = '/db/betabarrel'
		GB.ortho = cli.ortho
		GB.indir = cli.output
		GB.cdd_on = cli.cdd
		GB.subject=cli.input
		GB.subject_gis = cli.subgi
		GB.target_gis = cli.targi
		GB.expect = cli.evalue
		GB.blast_all()
		print "Loading BLAST results"
		GB.load_good()
		print "Writing FASTAS"
		GB.write_fastas()
		print "Running HMMTOP"
		GB.hmmtop()
		print "Calculating TMS Overlap"
		GB.calculate_tms_scores()
		print "Writing Results"
		GB.write_results()
	else:
		opts.print_help()

